import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error
import warnings
from data_loader import EduDataLoader
import sys
from PySide6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QLabel, QLineEdit, QPushButton, QTableWidget, QTableWidgetItem,
    QTabWidget, QHeaderView
)
from PySide6.QtCore import Qt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
import matplotlib.pyplot as plt

warnings.filterwarnings('ignore')

class StudentGradePredictor:
    def __init__(self, data_loader):
        self.data_loader = data_loader
        self.models = {}  # 各科目预测模型
        self.current_grades = None
        self.attendance_impact = None

    def prepare_data(self):
        """准备训练数据"""
        # 加载所有数据
        data = self.data_loader.load_all()

        # 合并成绩和考勤数据
        scores = data['scores']
        attendance = data['attendance']

        # 计算每个学生的迟到早退次数
        attendance_counts = attendance.groupby('bf_studentID')['severity'].agg(['count', 'sum']).reset_index()
        attendance_counts.columns = ['bf_studentID', 'total_attendance_records', 'total_severity']

        # 计算当前各科成绩（取最近一次考试）
        self.current_grades = scores.sort_values('exam_date').groupby(['bf_studentID', 'subject']).last()[
            'score'].unstack()

        # 合并特征
        features = pd.merge(self.current_grades.reset_index(),
                            attendance_counts,
                            on='bf_studentID',
                            how='left').fillna(0)

        # 添加学生基本信息
        students = data['students'][['bf_StudentID', 'bf_sex', 'age', 'cla_Name']]
        features = pd.merge(features, students,
                            left_on='bf_studentID',
                            right_on='bf_StudentID',
                            how='left').drop('bf_StudentID', axis=1)

        # 计算考勤影响因子（迟到早退越多，影响越大）
        self.attendance_impact = features[['bf_studentID', 'total_severity']].set_index('bf_studentID')

        return features

    def train_models(self):
        """训练各科目预测模型"""
        data = self.prepare_data()
        subjects = [col for col in data.columns if col not in ['bf_studentID', 'total_attendance_records',
                                                               'total_severity', 'bf_sex', 'age', 'cla_Name']]

        for subject in subjects:
            # 准备训练数据
            subject_data = data.dropna(subset=[subject])
            X = subject_data[['total_severity', 'bf_sex', 'age']]  # 使用考勤严重程度、性别、年龄作为特征
            y = subject_data[subject]

            # 转换为数值
            X = pd.get_dummies(X, columns=['bf_sex'])

            # 分割训练测试集
            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

            # 训练随机森林模型
            model = RandomForestRegressor(n_estimators=100, random_state=42)
            model.fit(X_train, y_train)

            # 评估模型
            preds = model.predict(X_test)
            mse = mean_squared_error(y_test, preds)
            print(f"{subject}模型测试集MSE: {mse:.2f}")

            self.models[subject] = model

    def predict_grades(self, student_id):
        if not self.models:
            self.train_models()

        if student_id not in self.current_grades.index:
            return None, None

        # 获取当前成绩
        current = self.current_grades.loc[student_id]

        # 准备预测数据
        student_data = self.data_loader.load_students()
        student_info = student_data[student_data['bf_StudentID'] == student_id][['bf_sex', 'age']].iloc[0]
        severity = self.attendance_impact.loc[student_id, 'total_severity']

        # 创建特征DataFrame
        X_pred = pd.DataFrame({
            'total_severity': [severity],
            'bf_sex': [student_info['bf_sex']],
            'age': [student_info['age']]
        })
        X_pred = pd.get_dummies(X_pred, columns=['bf_sex'])

        # 预测各科成绩
        predicted = {}
        for subject, model in self.models.items():  # 直接遍历models字典
            # 确保特征一致
            for col in model.feature_names_in_:
                if col not in X_pred.columns:
                    X_pred[col] = 0

            # 按模型要求的特征顺序排列
            X_pred = X_pred[model.feature_names_in_]
            predicted[subject] = model.predict(X_pred)[0]

        # 根据考勤调整预测成绩
        adjustment_factor = 1 - (severity / (severity + 10)) * 0.2
        predicted = pd.Series(predicted) * adjustment_factor

        return current, predicted.round(1)

# 新增可视化界面类
class GradePredictorUI(QMainWindow):
    def __init__(self, predictor):
        super().__init__()
        self.predictor = predictor
        self.initUI()

    def initUI(self):
        self.setWindowTitle('学生成绩预测系统')
        self.setGeometry(300, 300, 1000, 700)

        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout()

        # 输入区域
        input_layout = QHBoxLayout()
        input_layout.addWidget(QLabel('学生学号:'))
        self.student_id_input = QLineEdit()
        self.predict_btn = QPushButton('预测成绩')
        self.predict_btn.clicked.connect(self.predict_grades)
        input_layout.addWidget(self.student_id_input)
        input_layout.addWidget(self.predict_btn)
        layout.addLayout(input_layout)

        # 结果显示区域
        self.result_tabs = QTabWidget()

        # 成绩表格
        self.grade_table = QTableWidget()
        self.grade_table.setColumnCount(4)
        self.grade_table.setHorizontalHeaderLabels(['科目', '当前成绩', '预测成绩', '变化趋势'])
        self.grade_table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)

        # 可视化图表
        self.figure = plt.figure(figsize=(10, 6))
        self.canvas = FigureCanvas(self.figure)

        self.result_tabs.addTab(self.grade_table, "成绩详情")
        self.result_tabs.addTab(self.canvas, "趋势分析")
        layout.addWidget(self.result_tabs)

        central_widget.setLayout(layout)

    def predict_grades(self):
        try:
            student_id = int(self.student_id_input.text())
            current, predicted = self.predictor.predict_grades(student_id)

            if current is None:
                self.show_error("未找到该学生信息")
                return

            self.display_results(current, predicted)

        except ValueError:
            self.show_error("请输入有效的学号")

    def display_results(self, current, predicted):
        self.grade_table.setRowCount(0)
        subjects = current.index.tolist()

        # 填充表格数据
        for row, subject in enumerate(subjects):
            self.grade_table.insertRow(row)
            self.grade_table.setItem(row, 0, QTableWidgetItem(subject))
            self.grade_table.setItem(row, 1, QTableWidgetItem(str(current[subject])))
            self.grade_table.setItem(row, 2, QTableWidgetItem(str(predicted[subject])))

            change = predicted[subject] - current[subject]
            change_item = QTableWidgetItem(f"{change:+.1f}")
            change_item.setForeground(Qt.red if change < 0 else Qt.green)
            self.grade_table.setItem(row, 3, change_item)

        # 绘制对比图表
        self.figure.clear()
        ax = self.figure.add_subplot(111)

        x = range(len(subjects))
        width = 0.35

        ax.bar([i - width / 2 for i in x], current.values, width,
               label='当前成绩', color='#1f77b4')
        ax.bar([i + width / 2 for i in x], predicted[subjects].values, width,
               label='预测成绩', color='#ff7f0e')

        ax.set_xticks(x)
        ax.set_xticklabels(subjects, rotation=45, ha='right')
        ax.set_ylabel('分数')
        ax.set_title(f'学生 {self.student_id_input.text()} 成绩预测对比')
        ax.legend()
        self.figure.tight_layout()
        self.canvas.draw()

    def show_error(self, message):
        self.grade_table.setRowCount(1)
        self.grade_table.setItem(0, 0, QTableWidgetItem(message))
        self.figure.clear()
        self.canvas.draw()


# 使用示例
if __name__ == "__main__":
    # 初始化数据加载器
    app = QApplication(sys.argv)
    loader = EduDataLoader()
    predictor = StudentGradePredictor(loader)
    window = GradePredictorUI(predictor)
    window.show()
    sys.exit(app.exec_())